import random
from getPrimeNumber import get_prime_numbers
q = int(input("请输入素数q: "))  # 从控制台输入素数q
g = int(input("请输入原根α: "))  # 从控制台输入原根α

def random_int(num):
    return random.randint(1, num-1)

def calculate(g, a, q):
    '''
    计算公钥
    输入:参考书p202,g为α,a为私钥Xa,q为q,具体公式：(α^Xa) mod q
    输出:公钥Ya
    '''
    result = 1
    while a > 0:
        if a & 1 == 1:
            result = (result * g) % q
        g = (g * g) % q
        a >>= 1
    return result


A = random_int(q)
B = random_int(q)
# 测试书怕p202中用例
# A = 97
# B = 233
print("A获取的私钥Xa=", A)
print("B获取的私钥Xb=", B)

alpha = calculate(g, A, q)
bate = calculate(g, B, q)
print("A计算得到的alpha为：", alpha)
print("B计算得到的bate为：", bate)

print("------A与B交换alpha与bate------")
key_A = calculate(bate, A, q)
key_B = calculate(alpha, B, q)
print("A计算得到的key为：", key_A)
print("B计算得到的key为：", key_B)

'''
测试输出1：(书p202用例)
请输入素数q: 353 
请输入原根α: 3
A获取的私钥Xa= 97
B获取的私钥Xb= 233
A计算得到的alpha为： 40
B计算得到的bate为： 248
------A与B交换alpha与bate------
A计算得到的key为： 160
B计算得到的key为： 160
'''
# qList = get_prime_numbers(100, 255)
# gList = [2,5,2,6,3,3,2,3,2,2,6,5,2,5,2,2,2,19,5,2,3,2,3,2,6,3,7,7,6]
# for q,g in zip(qList,gList):
#     print("素数q=",q)
#     print("原根α=",g)
#     A = random_int(q)
#     B = random_int(q)
#     # 测试书怕p202中用例
#     # A = 97
#     # B = 233
#     print("A获取的私钥Xa=", A)
#     print("B获取的私钥Xb=", B)

#     alpha = calculate(g, A, q)
#     bate = calculate(g, B, q)
#     print("A计算得到的alpha为：", alpha)
#     print("B计算得到的bate为：", bate)

#     print("------A与B交换alpha与bate------")
#     key_A = calculate(bate, A, q)
#     key_B = calculate(alpha, B, q)
#     print("A计算得到的key为：", key_A)
#     print("B计算得到的key为：", key_B)
#     print("")

'''
测试输出2：(原根和素数按照ppt表中要求，私钥随机生成)
素数q= 101
原根α= 2
A获取的私钥Xa= 83
B获取的私钥Xb= 90
A计算得到的alpha为： 66
B计算得到的bate为： 65
------A与B交换alpha与bate------
A计算得到的key为： 6
B计算得到的key为： 6

素数q= 103
原根α= 5
A获取的私钥Xa= 49
B获取的私钥Xb= 100
A计算得到的alpha为： 70
B计算得到的bate为： 33
------A与B交换alpha与bate------
A计算得到的key为： 7
B计算得到的key为： 7

素数q= 107
原根α= 2
A获取的私钥Xa= 43
B获取的私钥Xb= 86
A计算得到的alpha为： 7
B计算得到的bate为： 49
------A与B交换alpha与bate------
A计算得到的key为： 25
B计算得到的key为： 25

素数q= 109
原根α= 6
A获取的私钥Xa= 48
B获取的私钥Xb= 52
A计算得到的alpha为： 27
B计算得到的bate为： 3
------A与B交换alpha与bate------
A计算得到的key为： 16
B计算得到的key为： 16

素数q= 113
原根α= 3
A获取的私钥Xa= 31
B获取的私钥Xb= 41
A计算得到的alpha为： 47
B计算得到的bate为： 23
------A与B交换alpha与bate------
A计算得到的key为： 103
B计算得到的key为： 103

素数q= 127
原根α= 3
A获取的私钥Xa= 54
B获取的私钥Xb= 101
A计算得到的alpha为： 64
B计算得到的bate为： 110
------A与B交换alpha与bate------
A计算得到的key为： 16
B计算得到的key为： 16

素数q= 131
原根α= 2
A获取的私钥Xa= 108
B获取的私钥Xb= 67
A计算得到的alpha为： 114
B计算得到的bate为： 127
------A与B交换alpha与bate------
A计算得到的key为： 27
B计算得到的key为： 27

素数q= 137
原根α= 3
A获取的私钥Xa= 38
B获取的私钥Xb= 122
A计算得到的alpha为： 17
B计算得到的bate为： 11
------A与B交换alpha与bate------
A计算得到的key为： 18
B计算得到的key为： 18

素数q= 139
原根α= 2
A获取的私钥Xa= 50
B获取的私钥Xb= 71
A计算得到的alpha为： 7
B计算得到的bate为： 135
------A与B交换alpha与bate------
A计算得到的key为： 49
B计算得到的key为： 49

素数q= 149
原根α= 2
A获取的私钥Xa= 124
B获取的私钥Xb= 25
A计算得到的alpha为： 17
B计算得到的bate为： 79
------A与B交换alpha与bate------
A计算得到的key为： 39
B计算得到的key为： 39

素数q= 151
原根α= 6
A获取的私钥Xa= 25
B获取的私钥Xb= 130
A计算得到的alpha为： 33
B计算得到的bate为： 16
------A与B交换alpha与bate------
A计算得到的key为： 118
B计算得到的key为： 118

素数q= 157
原根α= 5
A获取的私钥Xa= 70
B获取的私钥Xb= 27
A计算得到的alpha为： 122
B计算得到的bate为： 65
------A与B交换alpha与bate------
A计算得到的key为： 141
B计算得到的key为： 141

素数q= 163
原根α= 2
A获取的私钥Xa= 115
B获取的私钥Xb= 1
A计算得到的alpha为： 89
B计算得到的bate为： 2
------A与B交换alpha与bate------
A计算得到的key为： 89
B计算得到的key为： 89

素数q= 167
原根α= 5
A获取的私钥Xa= 156
B获取的私钥Xb= 61
A计算得到的alpha为： 54
B计算得到的bate为： 37
------A与B交换alpha与bate------
A计算得到的key为： 85
B计算得到的key为： 85

素数q= 173
原根α= 2
A获取的私钥Xa= 83
B获取的私钥Xb= 124
A计算得到的alpha为： 108
B计算得到的bate为： 84
------A与B交换alpha与bate------
A计算得到的key为： 29
B计算得到的key为： 29

素数q= 179
原根α= 2
A获取的私钥Xa= 138
B获取的私钥Xb= 3
A计算得到的alpha为： 5
B计算得到的bate为： 8
------A与B交换alpha与bate------
A计算得到的key为： 125
B计算得到的key为： 125

素数q= 181
原根α= 2
A获取的私钥Xa= 19
B获取的私钥Xb= 97
A计算得到的alpha为： 112
B计算得到的bate为： 53
------A与B交换alpha与bate------
A计算得到的key为： 131
B计算得到的key为： 131

素数q= 191
原根α= 19
A获取的私钥Xa= 140
B获取的私钥Xb= 46
A计算得到的alpha为： 153
B计算得到的bate为： 149
------A与B交换alpha与bate------
A计算得到的key为： 121
B计算得到的key为： 121

素数q= 193
原根α= 5
A获取的私钥Xa= 85
B获取的私钥Xb= 110
A计算得到的alpha为： 15
B计算得到的bate为： 137
------A与B交换alpha与bate------
A计算得到的key为： 101
B计算得到的key为： 101

素数q= 197
原根α= 2
A获取的私钥Xa= 53
B获取的私钥Xb= 81
A计算得到的alpha为： 170
B计算得到的bate为： 147
------A与B交换alpha与bate------
A计算得到的key为： 111
B计算得到的key为： 111

素数q= 199
原根α= 3
A获取的私钥Xa= 36
B获取的私钥Xb= 116
A计算得到的alpha为： 103
B计算得到的bate为： 91
------A与B交换alpha与bate------
A计算得到的key为： 125
B计算得到的key为： 125

素数q= 211
原根α= 2
A获取的私钥Xa= 128
B获取的私钥Xb= 161
A计算得到的alpha为： 119
B计算得到的bate为： 111
------A与B交换alpha与bate------
A计算得到的key为： 201
B计算得到的key为： 201

素数q= 223
原根α= 3
A获取的私钥Xa= 176
B获取的私钥Xb= 88
A计算得到的alpha为： 201
B计算得到的bate为： 76
------A与B交换alpha与bate------
A计算得到的key为： 126
B计算得到的key为： 126

素数q= 227
原根α= 2
A获取的私钥Xa= 93
B获取的私钥Xb= 213
A计算得到的alpha为： 18
B计算得到的bate为： 193
------A与B交换alpha与bate------
A计算得到的key为： 204
B计算得到的key为： 204

素数q= 229
原根α= 6
A获取的私钥Xa= 225
B获取的私钥Xb= 180
A计算得到的alpha为： 88
B计算得到的bate为： 104
------A与B交换alpha与bate------
A计算得到的key为： 43
B计算得到的key为： 43

素数q= 233
原根α= 3
A获取的私钥Xa= 64
B获取的私钥Xb= 102
A计算得到的alpha为： 126
B计算得到的bate为： 161
------A与B交换alpha与bate------
A计算得到的key为： 142
B计算得到的key为： 142

素数q= 239
原根α= 7
A获取的私钥Xa= 162
B获取的私钥Xb= 233
A计算得到的alpha为： 226
B计算得到的bate为： 149
------A与B交换alpha与bate------
A计算得到的key为： 55
B计算得到的key为： 55

素数q= 241
原根α= 7
A计算得到的alpha为： 234
B计算得到的bate为： 198
------A与B交换alpha与bate------
A计算得到的key为： 43
B计算得到的key为： 43

素数q= 251
原根α= 6
A获取的私钥Xa= 26
B获取的私钥Xb= 90
A计算得到的alpha为： 75
B计算得到的bate为： 51
------A与B交换alpha与bate------
A计算得到的key为： 51
B计算得到的key为： 51

'''